import sys
import cv2
import os
import numpy as np
import pickle
from datetime import datetime
from PyQt5.QtWidgets import *
from PyQt5.QtGui import *
from PyQt5.QtCore import *


class FaceRecognitionSystem(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowIcon(QIcon("icon.png"))  # 设置窗口图标
        self.cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        self.face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
        self.face_recognizer = cv2.face.LBPHFaceRecognizer_create()
        self.label_map = {}
        self.frame_counter = 0
        self.is_paused = False
        self.conf_threshold = 80

        # 加载背景图
        self.background_image = QImage("img/3.png")
        if self.background_image.isNull():
            # 如果背景图加载失败，使用默认颜色
            self.setStyleSheet("background-color: blue;")
        else:
            self.setAutoFillBackground(True)
            palette = self.palette()
            brush = QBrush(self.background_image.scaled(
                self.size(), Qt.IgnoreAspectRatio, Qt.SmoothTransformation))
            palette.setBrush(QPalette.Window, brush)
            self.setPalette(palette)

        self.init_ui()
        self.load_dataset()
        self.timer = self.startTimer(30)

    def init_ui(self):
        self.setWindowTitle("智能人脸识别考勤系统")
        self.setGeometry(100, 100, 1200, 700)

        # 设置整体样式
        self.setStyleSheet("""
            QWidget {
                font-family: "Microsoft YaHei", "Arial", sans-serif;
            }
            QLabel {
                color: #e2e8f0;
                font-size: 14px;
            }
            QLineEdit {
                background-color: rgba(255, 255, 255, 0.1);
                border: 1px solid rgba(255, 255, 255, 0.3);
                border-radius: 5px;
                padding: 8px;
                color: white;
                font-size: 14px;
            }
            QLineEdit::placeholder {
                color: rgba(255, 255, 255, 0.5);
            }
            QPushButton {
                background: linear-gradient(to right, #4facfe 0%, #00f2fe 100%);
                color: white;
                border-radius: 8px;
                padding: 10px 20px;
                font-size: 14px;
                font-weight: 500;
                border: none;
                transition: all 0.3s ease;
            }
            QPushButton:hover {
                background: linear-gradient(to right, #00f2fe 0%, #4facfe 100%);
                transform: translateY(-2px);
                box-shadow: 0 4px 12px rgba(0, 242, 254, 0.3);
            }
            QPushButton:pressed {
                transform: translateY(1px);
                box-shadow: 0 2px 6px rgba(0, 242, 254, 0.3);
            }
            QGroupBox {
                background-color: rgba(255, 255, 255, 0.05);
                border: 1px solid rgba(255, 255, 255, 0.1);
                border-radius: 10px;
                padding: 10px;
                margin-top: 20px;
            }
            QGroupBox::title {
                color: #00f2fe;
                subcontrol-position: top center;
                padding: 0 10px;
                font-size: 16px;
                font-weight: bold;
            }
            QSlider::groove:horizontal {
                border: 1px solid #999999;
                height: 8px;
                background: rgba(255, 255, 255, 0.1);
                margin: 2px 0;
                border-radius: 4px;
            }
            QSlider::handle:horizontal {
                background: qlineargradient(x1:0, y1:0, x2:1, y2:1, stop:0 #4facfe, stop:1 #00f2fe);
                border: 1px solid #5c5c5c;
                width: 18px;
                margin: -4px 0;
                border-radius: 9px;
            }
        """)

        main_layout = QHBoxLayout()
        main_layout.setSpacing(20)
        main_layout.setContentsMargins(20, 20, 20, 20)

        # 左侧视频显示区
        video_panel = QWidget()
        video_panel.setStyleSheet("background-color: rgba(255, 255, 255, 0.05); border-radius: 15px;")
        video_layout = QVBoxLayout()
        video_layout.setContentsMargins(15, 15, 15, 15)

        # 标题
        title_label = QLabel("实时人脸识别")
        title_label.setAlignment(Qt.AlignCenter)
        title_label.setStyleSheet("color: #00f2fe; font-size: 20px; font-weight: bold; margin-bottom: 10px;")
        video_layout.addWidget(title_label)

        self.video_label = QLabel()
        self.video_label.setFixedSize(800, 550)
        self.video_label.setStyleSheet("background-color: rgba(0, 0, 0, 0.5); border-radius: 10px;")
        self.video_label.setAlignment(Qt.AlignCenter)

        # 添加加载动画容器
        self.loading_placeholder = QLabel("正在加载摄像头...")
        self.loading_placeholder.setAlignment(Qt.AlignCenter)
        self.loading_placeholder.setStyleSheet("color: #e2e8f0; font-size: 16px;")
        video_layout.addWidget(self.loading_placeholder)
        video_layout.addWidget(self.video_label)

        video_panel.setLayout(video_layout)
        main_layout.addWidget(video_panel, 7)

        # 右侧控制面板
        control_panel = QWidget()
        control_panel.setStyleSheet("background-color: rgba(255, 255, 255, 0.05); border-radius: 15px;")
        control_layout = QVBoxLayout()
        control_layout.setSpacing(20)
        control_layout.setContentsMargins(20, 20, 20, 20)

        # 系统状态显示
        status_frame = QFrame()
        status_frame.setStyleSheet("background-color: rgba(255, 255, 255, 0.05); border-radius: 10px; padding: 15px;")
        status_layout = QVBoxLayout()

        self.status_label = QLabel("就绪 | 摄像头正常")
        self.status_label.setStyleSheet("color: #00f2fe; font-size: 16px; font-weight: bold;")
        status_layout.addWidget(self.status_label)

        # 识别进度条
        progress_frame = QFrame()
        progress_frame.setStyleSheet("margin-top: 10px;")
        progress_layout = QHBoxLayout()
        progress_layout.setContentsMargins(0, 0, 0, 0)

        progress_label = QLabel("识别进度:")
        progress_label.setStyleSheet("color: #e2e8f0;")
        self.progress_bar = QProgressBar()
        self.progress_bar.setRange(0, 100)
        self.progress_bar.setValue(0)
        self.progress_bar.setStyleSheet("""
            QProgressBar {
                border: 1px solid rgba(255, 255, 255, 0.2);
                border-radius: 5px;
                text-align: center;
                background-color: rgba(255, 255, 255, 0.1);
                height: 20px;
            }
            QProgressBar::chunk {
                background: qlineargradient(x1:0, y1:0, x2:1, y2:0, stop:0 #4facfe, stop:1 #00f2fe);
                border-radius: 4px;
            }
        """)

        progress_layout.addWidget(progress_label)
        progress_layout.addWidget(self.progress_bar)
        progress_frame.setLayout(progress_layout)

        status_layout.addWidget(progress_frame)
        status_frame.setLayout(status_layout)
        control_layout.addWidget(status_frame)

        # 输入框组
        input_group = QGroupBox("用户信息")
        input_layout = QFormLayout()
        input_layout.setVerticalSpacing(15)

        name_label = QLabel("姓名:")
        name_label.setStyleSheet("color: #e2e8f0;")
        self.name_edit = QLineEdit()
        self.name_edit.setPlaceholderText("请输入姓名")

        id_label = QLabel("学号:")
        id_label.setStyleSheet("color: #e2e8f0;")
        self.id_edit = QLineEdit()
        self.id_edit.setPlaceholderText("请输入学号")

        input_layout.addRow(name_label, self.name_edit)
        input_layout.addRow(id_label, self.id_edit)
        input_group.setLayout(input_layout)
        control_layout.addWidget(input_group)

        # 控制按钮组
        btn_group = QGroupBox("操作")
        btn_layout = QVBoxLayout()
        btn_layout.setSpacing(15)

        self.play_pause_btn = QPushButton("开始识别")
        self.play_pause_btn.setIcon(QIcon.fromTheme("media-playback-start"))
        self.play_pause_btn.setMinimumHeight(40)
        self.play_pause_btn.clicked.connect(self.toggle_play_pause)

        self.capture_btn = QPushButton("采集人脸")
        self.capture_btn.setIcon(QIcon.fromTheme("camera-photo"))
        self.capture_btn.setMinimumHeight(40)
        self.capture_btn.clicked.connect(self.capture_face)

        self.punch_btn = QPushButton("立即打卡")
        self.punch_btn.setIcon(QIcon.fromTheme("system-run"))
        self.punch_btn.setMinimumHeight(40)
        self.punch_btn.clicked.connect(self.punch_in)

        btn_layout.addWidget(self.play_pause_btn)
        btn_layout.addWidget(self.capture_btn)
        btn_layout.addWidget(self.punch_btn)
        btn_group.setLayout(btn_layout)
        control_layout.addWidget(btn_group)

        # 识别灵敏度设置
        sensitivity_group = QGroupBox("识别灵敏度")
        sensitivity_layout = QVBoxLayout()

        slider_layout = QHBoxLayout()
        slider_label = QLabel("阈值:")
        slider_label.setStyleSheet("color: #e2e8f0;")
        self.threshold_label = QLabel(f"{self.conf_threshold}%")
        self.threshold_label.setStyleSheet("color: #00f2fe; font-weight: bold;")

        self.conf_slider = QSlider(Qt.Horizontal)
        self.conf_slider.setRange(60, 95)
        self.conf_slider.setValue(self.conf_threshold)
        self.conf_slider.setMinimumHeight(30)
        self.conf_slider.valueChanged.connect(self.update_threshold)

        slider_layout.addWidget(slider_label)
        slider_layout.addWidget(self.conf_slider)
        slider_layout.addWidget(self.threshold_label)

        sensitivity_layout.addLayout(slider_layout)
        sensitivity_group.setLayout(sensitivity_layout)
        control_layout.addWidget(sensitivity_group)

        # 考勤记录查看按钮
        record_btn = QPushButton("查看考勤记录")
        record_btn.setStyleSheet("""
            QPushButton {
                background: linear-gradient(to right, #6a11cb 0%, #2575fc 100%);
                color: white;
                border-radius: 8px;
                padding: 10px 20px;
                font-size: 14px;
                font-weight: 500;
                border: none;
                transition: all 0.3s ease;
            }
            QPushButton:hover {
                background: linear-gradient(to right, #2575fc 0%, #6a11cb 100%);
                transform: translateY(-2px);
                box-shadow: 0 4px 12px rgba(106, 17, 203, 0.3);
            }
            QPushButton:pressed {
                transform: translateY(1px);
                box-shadow: 0 2px 6px rgba(106, 17, 203, 0.3);
            }
        """)
        record_btn.setMinimumHeight(40)
        record_btn.clicked.connect(self.show_records)
        control_layout.addWidget(record_btn)

        # 底部信息
        footer_label = QLabel("智能人脸识别考勤系统 v1.0")
        footer_label.setAlignment(Qt.AlignCenter)
        footer_label.setStyleSheet("color: rgba(255, 255, 255, 0.5); font-size: 12px; margin-top: 10px;")
        control_layout.addWidget(footer_label)

        control_panel.setLayout(control_layout)
        main_layout.addWidget(control_panel, 3)

        self.setLayout(main_layout)

    def update_frame(self):
        if self.is_paused:
            return

        ret, frame = self.cap.read()
        if not ret:
            self.status_label.setText("错误：无法访问摄像头！")
            self.status_label.setStyleSheet("color: #e74c3c; font-weight: bold;")
            self.loading_placeholder.show()
            self.video_label.hide()
            return

        # 摄像头成功加载后隐藏加载提示
        self.loading_placeholder.hide()
        self.video_label.show()

        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = self.face_cascade.detectMultiScale(gray, 1.3, 5)

        # 更新进度条
        self.progress_bar.setValue(len(faces) * 25 if len(faces) < 4 else 100)

        for (x, y, w, h) in faces:
            # 绘制彩色边框
            cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 255), 2)
            if self.play_pause_btn.text() == "停止识别":
                face_roi = gray[y:y + h, x:x + w]
                try:
                    label, conf = self.face_recognizer.predict(face_roi)
                    if conf < self.conf_threshold:
                        user_id = self.label_map.get(label, "未知")
                        confidence_percent = 100 - conf

                        # 绘制识别信息
                        cv2.rectangle(frame, (x, y - 30), (x + w, y), (0, 255, 255), -1)
                        cv2.putText(frame, f"{user_id} ({confidence_percent:.1f}%)",
                                    (x + 5, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 0), 2)

                        # 更新状态信息
                        self.status_label.setText(f"已识别: {user_id} | 置信度: {confidence_percent:.1f}%")
                        self.status_label.setStyleSheet("color: #00f2fe; font-size: 16px; font-weight: bold;")
                except:
                    pass

        # 转换图像格式
        h, w = frame.shape[:2]
        q_img = QImage(frame.data, w, h, 3 * w, QImage.Format_BGR888)
        self.video_label.setPixmap(QPixmap.fromImage(q_img).scaled(
            self.video_label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation))

    def toggle_play_pause(self):
        self.is_paused = not self.is_paused
        if self.is_paused:
            self.play_pause_btn.setText("开始识别")
            self.play_pause_btn.setIcon(QIcon.fromTheme("media-playback-start"))
            self.status_label.setText("识别已暂停")
            self.status_label.setStyleSheet("color: #e2e8f0; font-size: 16px; font-weight: bold;")
        else:
            self.play_pause_btn.setText("停止识别")
            self.play_pause_btn.setIcon(QIcon.fromTheme("media-playback-stop"))
            self.status_label.setText("正在识别...")
            self.status_label.setStyleSheet("color: #00f2fe; font-size: 16px; font-weight: bold;")

    def update_threshold(self, value):
        self.conf_threshold = value
        self.threshold_label.setText(f"{value}%")
        self.status_label.setText(f"识别灵敏度已更新为: {value}%")

    def load_dataset(self):
        images = []
        labels = []
        label_map = {}
        current_label = 0

        if not os.path.exists('pic'):
            os.makedirs('pic')

        for filename in os.listdir('pic'):
            if filename.endswith('.jpg'):
                id_num = filename.split('-')[0]
                if id_num not in label_map:
                    label_map[id_num] = current_label
                    current_label += 1
                label = label_map[id_num]

                img_path = os.path.join('pic', filename)
                img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
                if img is not None:
                    img = cv2.resize(img, (100, 100))
                    images.append(img)
                    labels.append(label)

        if len(images) > 0:
            self.face_recognizer.train(images, np.array(labels))
            self.face_recognizer.save("face_model.yml")
            with open("label_map.pkl", "wb") as f:
                pickle.dump(label_map, f)
            self.label_map = {v: k for k, v in label_map.items()}
            self.status_label.setText(f"已加载 {len(label_map)} 个用户数据")
        else:
            self.status_label.setText("未找到用户数据，请先采集人脸")
            self.status_label.setStyleSheet("color: #e74c3c; font-weight: bold;")

    def capture_face(self):
        name = self.name_edit.text().strip()
        id_num = self.id_edit.text().strip()

        if not id_num:
            self.show_message("输入错误", "学号不能为空！", QMessageBox.Warning)
            return

        ret, frame = self.cap.read()
        if not ret:
            self.show_message("摄像头错误", "无法读取摄像头数据。", QMessageBox.Critical)
            return

        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = self.face_cascade.detectMultiScale(gray, 1.3, 5)

        if len(faces) == 0:
            self.show_message("未检测到人脸", "请确保正对摄像头。", QMessageBox.Warning)
            return

        success_count = 0
        for i, (x, y, w, h) in enumerate(faces):
            face_img = gray[y:y + h, x:x + w]
            face_img = cv2.resize(face_img, (100, 100))
            timestamp = datetime.now().strftime("%Y%m%d%H%M%S%f")
            file_name = f"pic/{id_num}-{timestamp}.jpg"
            cv2.imwrite(file_name, face_img)
            success_count += 1

        self.load_dataset()
        self.show_message("采集成功", f"成功保存 {success_count} 张人脸图片！", QMessageBox.Information)

    def punch_in(self):
        if self.is_paused:
            self.show_message("系统暂停", "请先开始识别。", QMessageBox.Warning)
            return

        ret, frame = self.cap.read()
        if not ret:
            self.show_message("摄像头错误", "无法读取摄像头数据。", QMessageBox.Critical)
            return

        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = self.face_cascade.detectMultiScale(gray, 1.3, 5)

        if len(faces) == 0:
            self.show_message("未检测到人脸", "请确保正对摄像头。", QMessageBox.Warning)
            return

        for (x, y, w, h) in faces:
            face_roi = gray[y:y + h, x:x + w]
            try:
                label, conf = self.face_recognizer.predict(face_roi)
                if conf < self.conf_threshold:
                    user_id = self.label_map.get(label, "未知")
                    if user_id == "未知":
                        raise ValueError("未注册用户")
                    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    self.save_punch_record(user_id, current_time)

                    face_roi = cv2.resize(face_roi, (100, 100))
                    timestamp = datetime.now().strftime("%Y%m%d%H%M%S%f")
                    file_path = f"pic/{user_id}-{timestamp}.jpg"
                    cv2.imwrite(file_path, face_roi)
                    self.load_dataset()

                    os.makedirs("attendance_images", exist_ok=True)
                    img_filename = f"attendance_images/{user_id}_{current_time.replace(':', '-')}.jpg"
                    cv2.putText(frame, f"打卡时间：{current_time}", (10, 30),
                                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
                    cv2.imwrite(img_filename, frame)

                    self.show_message("打卡成功",
                                      f"{user_id} 于 {current_time} 打卡成功！\n"
                                      f"图像已保存：{img_filename}",
                                      QMessageBox.Information)
                    return
                else:
                    raise ValueError("识别失败")
            except:
                self.show_message("识别失败", "未找到匹配的用户，请重新尝试。", QMessageBox.Warning)
                return

    def save_punch_record(self, user_id, time_str):
        with open("attendance_records.csv", "a", encoding="utf-8-sig") as f:
            f.write(f"{user_id},{time_str}\n")

    def show_records(self):
        """显示考勤记录"""
        if not os.path.exists("attendance_records.csv"):
            self.show_message("记录为空", "暂无考勤记录。", QMessageBox.Information)
            return

        try:
            with open("attendance_records.csv", "r", encoding="utf-8-sig") as f:
                records = f.readlines()

            if not records:
                self.show_message("记录为空", "暂无考勤记录。", QMessageBox.Information)
                return

            # 创建记录窗口
            record_window = QDialog(self)
            record_window.setWindowTitle("考勤记录")
            record_window.setGeometry(200, 200, 600, 400)
            record_window.setStyleSheet("background-color: #1a1a2e;")

            layout = QVBoxLayout()

            # 添加表格
            table = QTableWidget(len(records), 2)
            table.setHorizontalHeaderLabels(["学号", "打卡时间"])
            table.setStyleSheet("""
                QTableWidget {
                    gridline-color: rgba(255, 255, 255, 0.1);
                    selection-background-color: #4facfe;
                }
                QTableWidget::item {
                    color: white;
                    padding: 5px;
                }
                QHeaderView::section {
                    background-color: rgba(255, 255, 255, 0.05);
                    color: #00f2fe;
                    padding: 5px;
                    border: 1px solid rgba(255, 255, 255, 0.1);
                }
            """)

            # 填充表格
            for i, record in enumerate(records):
                user_id, time_str = record.strip().split(',')
                table.setItem(i, 0, QTableWidgetItem(user_id))
                table.setItem(i, 1, QTableWidgetItem(time_str))

            table.resizeColumnsToContents()
            layout.addWidget(table)

            # 添加关闭按钮
            close_btn = QPushButton("关闭")
            close_btn.setStyleSheet("""
                QPushButton {
                    background: linear-gradient(to right, #4facfe 0%, #00f2fe 100%);
                    color: white;
                    border-radius: 5px;
                    padding: 8px 20px;
                    font-size: 14px;
                    border: none;
                }
                QPushButton:hover {
                    background: linear-gradient(to right, #00f2fe 0%, #4facfe 100%);
                }
            """)
            close_btn.clicked.connect(record_window.close)
            layout.addWidget(close_btn, alignment=Qt.AlignCenter)

            record_window.setLayout(layout)
            record_window.exec_()

        except Exception as e:
            self.show_message("错误", f"读取记录失败: {str(e)}", QMessageBox.Critical)

    def show_message(self, title, message, icon):
        """显示消息框的统一方法"""
        msg_box = QMessageBox(self)
        msg_box.setWindowTitle(title)
        msg_box.setText(message)
        msg_box.setIcon(icon)
        msg_box.setStandardButtons(QMessageBox.Ok)

        # 自定义消息框样式
        msg_box.setStyleSheet("""
            QMessageBox {
                background-color: #1a1a2e;
                border: 1px solid rgba(255, 255, 255, 0.2);
                border-radius: 10px;
            }
            QLabel {
                color: white;
                font-size: 14px;
            }
            QPushButton {
                background: linear-gradient(to right, #4facfe 0%, #00f2fe 100%);
                color: white;
                border-radius: 5px;
                padding: 8px 20px;
                font-size: 14px;
                border: none;
                min-width: 80px;
            }
            QPushButton:hover {
                background: linear-gradient(to right, #00f2fe 0%, #4facfe 100%);
            }
        """)

        msg_box.exec_()

    def timerEvent(self, event):
        self.update_frame()

    def resizeEvent(self, event):
        # 调整背景图大小
        if not self.background_image.isNull():
            palette = self.palette()
            brush = QBrush(self.background_image.scaled(
                self.size(), Qt.IgnoreAspectRatio, Qt.SmoothTransformation))
            palette.setBrush(QPalette.Window, brush)
            self.setPalette(palette)
        super().resizeEvent(event)

    def closeEvent(self, event):
        self.cap.release()
        event.accept()


if __name__ == '__main__':
    app = QApplication(sys.argv)
    app.setStyle("Fusion")
    ex = FaceRecognitionSystem()
    ex.show()
    sys.exit(app.exec_())